/*
 * Copyright 2019 Chris Morgan <chmorgan@gmail.com>
 *
 * GPLv3 licensed, see LICENSE for full text
 * Commercial licensing available
 */

#define PERFORMANCE_CRITICAL_LOGGING

using System;
using System.Collections.Generic;

namespace PacketDotNet.Connections
{
    /// <summary>
    /// Attaches to a TcpFlow and reassembles the packets from this flow into
    /// a TcpStream
    /// </summary>
    public class TcpStreamGenerator : IDisposable
    {
        private static readonly log4net.ILog log = log4net.LogManager.GetLogger(System.Reflection.MethodBase.GetCurrentMethod().DeclaringType);        

        /// <summary>
        /// Depending on the callback mode used in the implementation, it must be
        /// able to handle some or all of the following conditions.  In a typical
        /// implementation, the callback is waiting for either NextInSequence or
        /// SizeLimitReached.  The DUPLICATE, OutOfSequence, OutOfRange would be
        /// ignored in most cases.
        /// </summary>
        public enum CallbackCondition
        {
            /// <summary>
            /// New packet was received and was in sequence
            /// </summary>
            NextInSequence,
            /// <summary>
            /// Size limit for the flow was reached
            /// </summary>
            SizeLimitReached,
            /// <summary>
            /// Duplicate packet was received and ignored
            /// </summary>
            DuplicateDropped,
            /// <summary>
            /// Packet was recieved out of sequence
            /// </summary>
            OutOfSequence,
            /// <summary>
            /// Packet was received out of the monitored range
            /// </summary>
            OutOfRange,
            /// <summary>
            /// Next packet in sequence was not received by 'timeout' amount of time
            /// </summary>
            ConnectionTimeout,
            /// <summary>
            /// A packet was recieved that overlaps with a previous packet but is a different size
            /// </summary>
            StreamError
        };

        /// <summary>
        /// Status returned by OnCallback that indicates to this class whether
        /// a given TcpFlow should continue to be monitored or whether monitoring should
        /// be stopped
        /// </summary>
        public enum CallbackReturnValue
        {
            /// <summary>
            /// Class should stop monitoring the flow
            /// </summary>
            StopMonitoring,

            /// <summary>
            /// Class should continue monitoring the flow
            /// </summary>
            ContinueMonitoring
        };

        /// <summary>
        /// Linked list that holds packets are are being reordered
        /// As packets are received they are inserted into this list into the appropriate
        /// position and then fed from the start of this list into the TcpStream until
        /// an out-of-order packet is encountered
        /// </summary>
        private LinkedList<TcpPacket> packets;

        /// <summary>
        /// The TcpStream generated by this class
        /// </summary>
        public TcpStream tcpStream;

        private CallbackCondition condition;


        private TcpFlow flow;
        private long? sizeLimitInBytes;

        // first sequence number to look for at start of monitoring
        private bool waiting_for_start_seq; // if true we haven't yet set start_seq
        private UInt32 start_seq;

        /// <summary>
        /// Sequence number of the last in-order packet
        /// </summary>
        public UInt32 LastContinuousSequence
        {
            get;
            private set;
        }

        /// <summary>
        /// Last overall sequence number
        /// </summary>
        public UInt32 last_overall_seq;

        private DateTime last_continuous_packet_received;

        /// <summary>
        /// If no packets are received for 'timeout' a ConnectionTimeout event is generated
        /// </summary>
        public TimeSpan timeout;

        /// <summary>
        /// Delegate for event updates from this class
        /// </summary>
        public delegate CallbackReturnValue CallbackDelegate(TcpStreamGenerator tcpStreamGenerator,
                                                             CallbackCondition condition);

        /// <summary>
        /// Called when this class has status to report such as new packets etc
        /// </summary>
        public event CallbackDelegate OnCallback;

        /// <summary>
        /// Create a TcpStreamGenerator given a flow
        /// </summary>
        /// <param name="f">
        /// A <see cref="TcpFlow"/>
        /// </param>
        /// <param name="timeout">
        /// A <see cref="TimeSpan"/>
        /// </param>
        /// <param name="sizeLimitInBytes">
        /// A <see cref="System.Nullable<System.Int64>"/>
        /// </param>
        public TcpStreamGenerator(TcpFlow f,
                                  TimeSpan timeout,
                                  long? sizeLimitInBytes)
        {
            this.flow = f;
            this.flow.OnPacketReceived += HandleFlowOnPacketReceived;
            this.timeout = timeout;
            this.sizeLimitInBytes = sizeLimitInBytes;

            packets = new LinkedList<PacketDotNet.TcpPacket>();

            tcpStream = new TcpStream();
            condition = TcpStreamGenerator.CallbackCondition.NextInSequence;
            waiting_for_start_seq = true;
            LastContinuousSequence = 0;
            last_overall_seq = 0;

            // use now as our last packet received time to avoid
            // timing this monitor out immediately
            last_continuous_packet_received = DateTime.Now;
        }

        void HandleFlowOnPacketReceived (SharpPcap.PosixTimeval timeval,
                                         TcpPacket tcp,
                                         TcpConnection connection,
                                         TcpFlow flow)
        {
            AddPacket(tcp);
        }

        /// <summary>
        /// Dispose
        /// </summary>
        public void Dispose ()
        {
            flow.OnPacketReceived -= HandleFlowOnPacketReceived;
        }

        /// <summary>
        /// Clear out all of the cached packets
        /// </summary>
        public void FreePackets()
        {
            packets.Clear();
        }

        /// <summary>
        /// returns true if the TcpStreamMonitor timed out
        /// </summary>
        /// <returns>
        /// A <see cref="System.Boolean"/>
        /// </returns>
        public bool CheckForTimeout()
        {
            TimeSpan ts = DateTime.Now - last_continuous_packet_received;
            if(ts.TotalSeconds > timeout.TotalSeconds)
            {
                // Monitor has timed out, notify the user of the monitor via the callback
                lock(this)
                {
                    condition = TcpStreamGenerator.CallbackCondition.ConnectionTimeout;
                }

                return true;
            }

            return false;
        }

        /// <summary>
        /// Handle a tcp packet
        /// </summary>
        /// <param name="newPacket">
        /// A <see cref="TcpPacket"/>
        /// </param>
        public void AddPacket(TcpPacket newPacket)
        {
            InternalAddPacket(newPacket);

            // report the state via a callback
            var retval = OnCallback(this, condition);

            if(retval == CallbackReturnValue.StopMonitoring)
            {
                Dispose();
            }
        }

        //FIXME: we don't handle rolling sequence numbers properly in here
        // so after 4GB of packets we will have issues
        private void InternalAddPacket(PacketDotNet.TcpPacket newPacket)
        {
#if PERFORMANCE_CRITICAL_LOGGING
            log.Debug("");
#endif

            UInt32 new_packet_seq = newPacket.SequenceNumber;

            // if we haven't set the value of start_seq we should do so now
            if(waiting_for_start_seq)
            {
#if PERFORMANCE_CRITICAL_LOGGING
                log.Debug("waiting_for_start_seq");
#endif
                // NOTE: we always use the sequence value because we won't have a new packet
                // added unless it matches the direction we are monitoring, see tcp_stream_manager::handle()
                start_seq = new_packet_seq;
                waiting_for_start_seq = false;
                LastContinuousSequence = last_overall_seq = new_packet_seq;
            }

            // Adjust sequence number against the start sequence number
#if PERFORMANCE_CRITICAL_LOGGING
            log.DebugFormat("new_packet.ip.tcp.get_seq() is {0}, start_seq is {1}",
                            newPacket.SequenceNumber, start_seq);
#endif

            UInt32 new_packet_size = (UInt32)newPacket.PayloadData.Length;

#if PERFORMANCE_CRITICAL_LOGGING
            log.DebugFormat("new_packet_seq: {0}, new_packet_size: {1}," +
                            " lastContinuousSequence: {2}, last_overall_seq: {3}",
                    new_packet_seq, new_packet_size,
                    LastContinuousSequence, last_overall_seq);
#endif

            // if the packet came before the last continuous sequence
            // number we've already seen the packet
            if(new_packet_seq < LastContinuousSequence)
            {
                log.DebugFormat("new_packet_seq({0}) < lastContinuousSequence({1}), rejecting packet",
                                new_packet_seq, LastContinuousSequence);
                return;
            }

            // Verify that the packet within our allowed sequence range
            if(SizeLimit.HasValue)
            {
                if ((new_packet_seq - start_seq) > SizeLimit)
                {
                    log.DebugFormat("OutOfRange, new_packet_seq({0}) > SizeLimit({1}), deleting packet",
                                    new_packet_seq, SizeLimit);
                    condition = CallbackCondition.OutOfRange;
                    return;
                }
            }

            bool packet_added = false;
            LinkedListNode<PacketDotNet.TcpPacket> pNode;

            // try to place this packet in our pending packets list
            if(packets.Count != 0)
            {
                pNode = packets.Last;
                do
                {
                    if(pNode == null)
                        break;

                    var p = pNode.Value;
                    UInt32 current_packet_seq = (UInt32)p.SequenceNumber;
                    UInt32 current_packet_size = (UInt32)p.PayloadData.Length;

#if PERFORMANCE_CRITICAL_LOGGING
                    log.DebugFormat("current_packet_seq: {0}, current_packet_size: {1}",
                                    current_packet_seq, current_packet_size);
#endif

                    // Does this packet belong after, but not *immediately* after current?
                    if (new_packet_seq > (current_packet_seq + current_packet_size))
                    {
#if PERFORMANCE_CRITICAL_LOGGING
                        log.Debug("OutOfSequence, adding packet");
#endif

                        // Sanity check passed, insert into vector
                        condition = CallbackCondition.OutOfSequence;
                        if (last_overall_seq < (new_packet_seq + new_packet_size))
                        {
                            last_overall_seq = new_packet_seq + new_packet_size;
                        }

                        packets.AddAfter(pNode, newPacket);
                        return; // adding an out-of-sequence packet won't put any other
                                // packets back in sequence, so we are done
                    }

                    // Does this packet fit exactly in sequence after current?
                    if (new_packet_seq == (current_packet_seq + current_packet_size))
                    {
#if PERFORMANCE_CRITICAL_LOGGING
                        log.Debug("packet fits into sequence, adding it");
#endif

                        // Verify highest sequence number
                        if ((new_packet_seq + new_packet_size) > last_overall_seq)
                            last_overall_seq = new_packet_seq + new_packet_size;

                        packets.AddAfter(pNode, newPacket);
                        packet_added = true;

                        break;
                    }

                    // if the sequence numbers match there is a packet overlap
                    if (new_packet_seq == current_packet_seq)
                    {
                        if (new_packet_size == current_packet_size)
                        {
#if PERFORMANCE_CRITICAL_LOGGING
                            log.Debug("DuplicateDropped");
#endif
                            condition = CallbackCondition.DuplicateDropped; // Duplicate packet, drop it
                            return;
                        } else
                        {
#if PERFORMANCE_CRITICAL_LOGGING
                            log.Debug("StreamError, packet would overlap previous packet and is different size");
#endif
                            condition = CallbackCondition.StreamError; // Error: Packet would overlap
                            return;
                        }
                    }

                    // move to the previous node
                    pNode = pNode.Previous;
                } while(packets.Count != 0);
            }

            // Packet did not fall between the existing packets or after all of them,
            // so it should belong before them all.
            if (!packet_added)
            {
#if PERFORMANCE_CRITICAL_LOGGING
                log.Debug("!packet_added, packet did not fall between existing packets or after them, adding packet at front");
#endif
                packets.AddFirst(newPacket);
                packet_added = true;
            }


            //
            // process our pending packets list
            // pull all sequential packets out of the packets list and
            // push them into the TcpStream
            //
            // retrieve the last packet
            // if there is no last packet then break out of the do {} while()
            pNode = packets.First;
            do
            {
                if(pNode == null)
                {
#if PERFORMANCE_CRITICAL_LOGGING
                    log.Debug("pNode == null");
#endif
                    break;
                }

                var p = pNode.Value;
                var seq = p.SequenceNumber;
                var size = (UInt32)p.PayloadData.Length;

                // if the size is zero and this is a syn packet then
                // we need to adjust the size to be 1 according to the tcp protocol
                if(newPacket.Synchronize)
                {
                    size = 1;
                    LastContinuousSequence = seq;

                    log.DebugFormat("syn packet, lastContinousSequence updated to {0}",
                                    LastContinuousSequence);
                }

#if PERFORMANCE_CRITICAL_LOGGING
                log.DebugFormat("seq: {0}, size: {1}", seq, size);
#endif

                if(seq == LastContinuousSequence)
                {
                    if(SizeLimit.HasValue)
                    {
                        // is this packet beyond our range?
                        if((seq - start_seq) >= SizeLimit + 1)
                        {
                            log.Debug("dropping beyond range packet");

                            // just drop this packet and loop
                            packets.Remove(pNode);
                            pNode = pNode.Next;
                            continue; // skip back around to the top of the do {} while()
                        }
                    }

                    condition = CallbackCondition.NextInSequence;

                    last_continuous_packet_received = DateTime.Now;

                    // store the packet and remove the node from the packets linked list
#if PERFORMANCE_CRITICAL_LOGGING
                    log.Debug("AppendPacket() to tcpStream");
#endif
                    tcpStream.AppendPacket(p);
                    packets.Remove(pNode);

                    // is this packet the next in line? if so update our
                    // last continuous sequence value
                    if(seq == LastContinuousSequence)
                    {
                        LastContinuousSequence += size;
                    }

                    // is the last continuous sequence after the current value?
                    // if our last overall sequence is behind, update it as well
                    if(last_overall_seq < LastContinuousSequence)
                    {
                        last_overall_seq = LastContinuousSequence;
                    }
                } else
                {
#if PERFORMANCE_CRITICAL_LOGGING
                    log.DebugFormat("seq {0}  != LastContinuousSequence() {1}",
                                    seq, LastContinuousSequence);
#endif
                    break;
                }

                pNode = pNode.Next;
            } while(packets.Count != 0);
        }

        /// <summary>
        /// Size limit for this stream generator, if one has been defined
        /// </summary>
        public long? SizeLimit
        {
            get { return sizeLimitInBytes; }
        }

        /// <summary>
        /// ToString() override
        /// </summary>
        /// <returns>
        /// A <see cref="System.String"/>
        /// </returns>
        public override string ToString ()
        {
            return string.Format("[TcpStreamGenerator: packets.Count: {0}, condition: {1}, tcpStream: {2}, waiting_for_start_seq: {3}, start_seq: {4}, LastContinuousSequence={5}, last_overall_seq: {6}, last_continuous_packet_received: {7}]",
                                 packets.Count,
                                 condition,
                                 tcpStream.ToString(),
                                 waiting_for_start_seq,
                                 start_seq,
                                 LastContinuousSequence,
                                 last_overall_seq,
                                 last_continuous_packet_received);
        }
    }
}
